#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版数据库连接模块 - 专为balance_rollback.py设计
"""

from sqlalchemy import create_engine


class db_connector:
    """
    简化数据库连接类，兼容原有代码调用方式
    """

    @staticmethod
    def conn_db(environment: str, db_type: str, db_name: str):
        """
        创建数据库连接引擎 - 兼容balance_rollback.py调用

        参数:
        environment: 环境标识（如'local'）
        db_type: 数据库类型（'production'或'test'）
        db_name: 数据库名称（'dwh_ods', 'dwh_dwd', 'dwh_ads', 'dwh_dcp'）

        返回:
        SQLAlchemy引擎对象
        """

        # 数据库连接配置
        connections = {
            # dolphischeduler环境配置
            ('ds', 'production', 'dwh_tmp'): {
                'host': '192.168.11.153',
                'port': 3306,
                'user': 'user_datahub',
                'password': 'bi_datahub',
                'database': 'dwh_tmp'
            },
            ('ds', 'test', 'dwh_tmp'): {
                'host': '192.168.186.152',
                'port': 3306,
                'user': 'user_datahub',
                'password': 'bi_datahub_test',
                'database': 'dwh_tmp'
            },
            # 正式环境配置
            ('local', 'production', 'dwh_tmp'): {
                'host': 'blj.oigit.com',
                'port': 33061,
                'user': 'b202d3dc-db1b-406e-9d59-31f651ac51c2',
                'password': 'AJa8RqwRX5EEQLg4',
                'database': 'dwh_tmp'
            },

            # 测试环境配置
            ('local', 'test', 'dwh_tmp'): {
                'host': '192.168.186.152',
                'port': 3306,
                'user': 'user_datahub',
                'password': 'bi_datahub_test',
                'database': 'dwh_tmp'
            }
        }

        # 获取连接配置
        key = (environment, db_type, db_name)
        config = connections.get(key)

        if not config:
            # 如果没有精确匹配，使用默认配置
            if db_type == 'production':
                config = connections[('local', 'production', 'dwh_ods')]
            else:
                config = connections[('local', 'test', 'dwh_ads')]

        # 创建连接字符串
        connection_string = (
            f"mysql+pymysql://{config['user']}:{config['password']}"
            f"@{config['host']}:{config['port']}/{config['database']}"
            f"?charset=utf8mb4"
        )

        # 创建并返回引擎
        return create_engine(connection_string)


# 测试函数
if __name__ == "__main__":
    print("测试数据库连接...")

    try:
        # 测试正式环境连接
        conn_prod = db_connector.conn_db('local', 'production', 'dwh_dwd')
        print("✓ 正式环境连接成功")

        # 测试测试环境连接
        conn_test = db_connector.conn_db('local', 'test', 'dwh_ads')
        print("✓ 测试环境连接成功")

    except Exception as e:
        print(f"✗ 连接失败: {str(e)}")